package org.framework.lazy.cloud.network.heartbeat.protocol;

import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.dns.*;
import lombok.extern.slf4j.Slf4j;

import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.HashMap;
import java.util.Map;

@Slf4j
public class DnsServerHandler extends SimpleChannelInboundHandler<DatagramDnsQuery> {


    private Map<String, String> domainIpMapping = new HashMap<>();


    {
        domainIpMapping.put("test.wu-framework.cn.", "124.222.152.160");
        domainIpMapping.put("www.google.com.hk.", "120.253.253.97");
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, DatagramDnsQuery query) throws Exception {

        // 创建DNS响应
        DatagramDnsResponse response = new DatagramDnsResponse(query.recipient(), query.sender(), query.id());
        DefaultDnsQuestion dnsQuestion = query.recordAt(DnsSection.QUESTION);
        response.addRecord(DnsSection.QUESTION, dnsQuestion);
        DnsRecordType dnsRecordType = dnsQuestion.type();

        log.info("dns:{}", dnsQuestion.name());
//        log.info("query:{}", query);
        String domain = dnsQuestion.name();
        if (dnsRecordType == DnsRecordType.A) {
            String ip = domainIpMapping.get(domain);
            if (ip != null) {
                log.info("ip:{}", ip);
                try {
                    InetAddress address = Inet4Address.getByName(ip);
                    DefaultDnsRawRecord answer = new DefaultDnsRawRecord(
                            domain, DnsRecordType.A, 3600,
                            Unpooled.wrappedBuffer(address.getAddress())
                    );
                    response.addRecord(DnsSection.ANSWER, answer);
                } catch (UnknownHostException e) {
                    e.printStackTrace();
                }

                // 发送响应
                ctx.writeAndFlush(response);
            } else {
                // 使用系统默认 DNS 解析
                try {
                    java.net.InetAddress[] addresses = java.net.InetAddress.getAllByName(domain);
                    for (java.net.InetAddress address : addresses) {
                        log.info("使用本地dns解析域名：{} ip:{}", domain, address.getHostAddress());
                        response.addRecord(DnsSection.ANSWER, new DefaultDnsRawRecord(
                                domain, DnsRecordType.A, 3600, Unpooled.wrappedBuffer(address.getAddress())));
                    }
//                    log.info("response:{}",response.toString());
                    ctx.writeAndFlush(response);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        } else {
            log.info("type:{}", dnsRecordType.name());
        }


    }


    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        cause.printStackTrace();
        ctx.close();
    }
}